
#include "token.h"

// Base64编码
int base64_encode(const char *input, char *output, size_t output_len) {
    BIO *bio, *b64;
    BUF_MEM *bptr;
    int ret = -1;

    b64 = BIO_new(BIO_f_base64());
    BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
    bio = BIO_new(BIO_s_mem());
    bio = BIO_push(b64, bio);

    // 写入数据
    BIO_write(bio, input, strlen(input));
    BIO_flush(bio);
    BIO_get_mem_ptr(bio, &bptr);

    // 检查输出缓冲区大小
    if (bptr->length < output_len) {
        memcpy(output, bptr->data, bptr->length);
        output[bptr->length] = '\0';
        ret = bptr->length;
    }

    BIO_free_all(bio);
    return ret;
}

// Base64解码
int base64_decode(const char *input, char *output, size_t output_len) {
    BIO *bio, *b64;
    int len = strlen(input);
    int ret = -1;
    
    // 输出缓冲区初始化为0
    memset(output, 0, output_len);
    
    // 创建BIO链
    bio = BIO_new_mem_buf(input, len);
    b64 = BIO_new(BIO_f_base64());
    BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
    bio = BIO_push(b64, bio);

    // 读取解码数据
    ret = BIO_read(bio, output, output_len - 1);
    if (ret > 0) output[ret] = '\0'; // 确保null终止
    
    BIO_free_all(bio);
    return ret;
}



// 生成JWT令牌
char* generate_jwt(const char *username) {
    // 创建Header
    json_t *header = json_object();
    json_object_set_new(header, "alg", json_string("HS256"));
    json_object_set_new(header, "typ", json_string("JWT"));
    char *header_str = json_dumps(header, JSON_COMPACT);
    json_decref(header);
    
    // 创建Payload
    time_t now = time(NULL);
    json_t *payload = json_object();
    json_object_set_new(payload, "sub", json_string(username));
    json_object_set_new(payload, "exp", json_integer(now + TOKEN_EXPIRY));
    char *payload_str = json_dumps(payload, JSON_COMPACT);
    json_decref(payload);
    
    // Base64编码Header和Payload
    char header_b64[256] = {0};
    char payload_b64[512] = {0};
    base64_encode(header_str, header_b64, sizeof(header_b64));
    base64_encode(payload_str, payload_b64, sizeof(payload_b64));
    
    free(header_str);
    free(payload_str);
    
    // 准备签名数据
    char signing_input[768] = {0};
    snprintf(signing_input, sizeof(signing_input), "%s.%s", header_b64, payload_b64);
    
    // HMAC-SHA256签名
    unsigned char digest[EVP_MAX_MD_SIZE];
    unsigned int digest_len;
    HMAC(EVP_sha256(), SECRET_KEY, strlen(SECRET_KEY),
        (unsigned char*)signing_input, strlen(signing_input),
        digest, &digest_len);
    
    // Base64编码签名
    char signature_b64[256] = {0};
    base64_encode((char*)digest, signature_b64, sizeof(signature_b64));
    
    // 组合完整Token
    static char token[1024] = {0};
    snprintf(token, sizeof(token), "%s.%s.%s", header_b64, payload_b64, signature_b64);
    return token;
}

// 验证JWT令牌（更健壮的处理）
int verify_jwt(const char *token, char *username, ssize_t username_size) {
    if (!token || !username) return -1;
    
    // 复制令牌以便修改
    char *token_copy = strdup(token);
    if (!token_copy) return -1;
    
    // 分割令牌部分
    char *parts[3] = {0};
    parts[0] = strtok(token_copy, ".");
    parts[1] = strtok(NULL, ".");
    parts[2] = strtok(NULL, ".");
    
    if (!parts[0] || !parts[1] || !parts[2]) {
        free(token_copy);
        return -2; // 无效令牌格式
    }
    
    // 准备验证签名数据
    char signing_input[768] = {0};
    snprintf(signing_input, sizeof(signing_input), "%s.%s", parts[0], parts[1]);
    
    // 解码签名
    char signature_bin[256] = {0};
    int sig_len = base64_decode(parts[2], signature_bin, sizeof(signature_bin));
    if (sig_len <= 0) {
        free(token_copy);
        return -3; // 无效签名
    }
    
    // 计算HMAC-SHA256签名
    unsigned char digest[EVP_MAX_MD_SIZE];
    unsigned int digest_len;
    HMAC(EVP_sha256(), SECRET_KEY, strlen(SECRET_KEY),
        (unsigned char*)signing_input, strlen(signing_input),
        digest, &digest_len);
    
    // 比较签名
    if (digest_len != (unsigned int)sig_len || 
        memcmp(digest, signature_bin, digest_len) != 0) {
        free(token_copy);
        return -4; // 签名不匹配
    }
    
    // 解码Payload
    char payload_json[512] = {0};
    int payload_len = base64_decode(parts[1], payload_json, sizeof(payload_json));
    if (payload_len <= 0) {
        free(token_copy);
        return -5; // 无效Payload
    }
    
    // 解析JSON
    json_error_t error;
    json_t *payload = json_loads(payload_json, 0, &error);
    if (!payload) {
        free(token_copy);
        fprintf(stderr, "JSON parse error: %s\n", error.text);
        return -6; // JSON解析失败
    }
    
    // 检查过期时间
    json_t *exp = json_object_get(payload, "exp");
    time_t exp_time = json_integer_value(exp);
    if (time(NULL) > exp_time) {
        json_decref(payload);
        free(token_copy);
        return -7; // 令牌过期
    }

    // 提取用户名
    json_t *sub = json_object_get(payload, "sub");
    const char *user = json_string_value(sub);
    if (!user) {
        json_decref(payload);
        free(token_copy);
        return -8; // 缺少用户字段
    }
    
    // 安全复制用户名
    strncpy(username, user, username_size - 1);
    username[username_size - 1] = '\0';
    
    // 清理资源
    json_decref(payload);
    free(token_copy);
    return 0; // 验证成功
}


// 添加在receive_request之后，命令处理之前
int verify_token_auth(int sockfd, protocol_header_t *header)
{
    // 登录相关命令不需要Token验证
    if (header->command == CMD_LOGIN_APPLY || header->command == CMD_LOGIN ||
        header->command == CMD_REGISTER || header->command == CMD_REGISTER_APPLY)
    {
        return 0;
    }
    printf("我收到的token是%s\n", header->token);
    // 空Token检查
    if (strlen(header->token) == 0)
    {
        char *msg = "未提供Token";
        sendCommand(sockfd, header->command, FAIL, msg, strlen(msg));
        return -1;
    }

    // 验证JWT Token
    char username[64];

    int ret = verify_jwt(header->token, username, sizeof(username));
    
    if (ret == 0) {
        printf("用户 %s 的Token验证成功\n", username);
        return 0;
    }


    // 错误处理
    char *msg = "";
    switch (ret)
    {
    case -1:
        msg = "无效Token格式";
        break;
    case -2:
        msg = "Token格式错误";
        break;
    case -3:
        msg = "Token签名验证失败";
        break;
    case -4:
        msg = "Token已过期";
        break;
    case -5:
        msg = "Token负载错误";
        break;
    default:
        msg = "未知Token错误";
    }

    printf("Token验证失败: %s\n", msg);
    sendCommand(sockfd, header->command, FAIL, msg, strlen(msg));
    return -2;
}